package cn.hiapi.core.basic.sharding.config;

import cn.hiapi.core.basic.sharding.algorithm.ShardingAlgorithmTable;
import cn.hiapi.core.basic.sharding.entity.BasicShardingEntity;
import cn.hiapi.core.basic.sharding.entity.ShardingOption;
import cn.hiapi.core.basic.sharding.enums.ShardingMode;
import cn.hiapi.core.basic.sharding.properties.ShardingModeProperties;
import cn.hiapi.core.basic.sharding.properties.ShardingTimeProperties;
import cn.hiapi.core.basic.sharding.utils.DatabaseUtils;
import cn.hiapi.core.basic.sharding.utils.ShardingUtils;
import com.alibaba.druid.pool.DruidDataSource;
import com.google.common.collect.Lists;
import org.apache.commons.lang3.StringUtils;
import org.apache.shardingsphere.driver.api.ShardingSphereDataSourceFactory;
import org.apache.shardingsphere.infra.config.algorithm.AlgorithmConfiguration;
import org.apache.shardingsphere.infra.config.mode.ModeConfiguration;
import org.apache.shardingsphere.mode.repository.standalone.StandalonePersistRepositoryConfiguration;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.strategy.keygen.KeyGenerateStrategyConfiguration;
import org.apache.shardingsphere.sharding.api.config.strategy.sharding.StandardShardingStrategyConfiguration;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import javax.sql.DataSource;
import java.util.*;
import java.util.stream.Collectors;

/**
 * @author AdinZ
 * @date 2023/12/15 17:48
 */
@Configuration
public class ShardingConfig {


    private Map<String, DataSource> initDataSource() {
        Map<String, DataSource> map = new HashMap<>();
        DruidDataSource dataSource = new DruidDataSource();
        dataSource.setDriverClassName("com.mysql.cj.jdbc.Driver");
        dataSource.setUrl("jdbc:mysql://127.0.0.1:3306/test?serverTimezone=UTC");
        dataSource.setUsername("root");
        dataSource.setPassword("adinz123");
        map.put("ds_1", dataSource);
        return map;
    }


    private ShardingRuleConfiguration shardingRuleConfiguration(Map<String, DataSource> dataSourceMap) throws Exception {
        ShardingRuleConfiguration configuration = new ShardingRuleConfiguration();

        ShardingTableScan tableScan = new ShardingTableScan(BasicShardingEntity.class, "cn/hiapi");
        buildTableConfig(dataSourceMap, tableScan.getTableClass(), configuration);
        configuration.setKeyGenerators(new HashMap<>() {{
            put("snowflake", new AlgorithmConfiguration("SNOWFLAKE", new Properties()));
        }});

        return configuration;
    }


    private void buildTableConfig(Map<String, DataSource> dataSourceMap, List<Class<?>> tables, ShardingRuleConfiguration configuration) {
        tables.forEach(e -> {
            try {
                BasicShardingEntity instance = ShardingTableScan.newInstance(e);
                ShardingOption option = instance.getOption();
                if (option.getMode() == ShardingMode.MODE) {
                    ShardingModeProperties properties = (ShardingModeProperties) option.getProperties();
                    String format = String.format("%s.%s_${0..%d}", option.getDatabase(), option.getTable(), (properties.getTableSize() - 1));
                    ShardingTableRuleConfiguration rule = new ShardingTableRuleConfiguration(option.getTable(), format);
                    final String algorithmName = option.getTable() + "_inline";
                    rule.setTableShardingStrategy(new StandardShardingStrategyConfiguration(option.getShardingColumn(), algorithmName));
                    if (StringUtils.isNotEmpty(option.getPrimaryKey()))
                        rule.setKeyGenerateStrategy(new KeyGenerateStrategyConfiguration(option.getPrimaryKey(), "snowflake"));
                    configuration.getTables().add(rule);
                    configuration.getShardingAlgorithms().put(algorithmName, new AlgorithmConfiguration("inline", new Properties() {{
                        put("algorithm-expression", option.getTable() + "_${" + option.getShardingColumn() + " % " + properties.getTableSize() + "}");
                    }}));
                } else if (option.getMode() == ShardingMode.TIME) {
                    ShardingTimeProperties properties = (ShardingTimeProperties) option.getProperties();
                    Set<String> logicTables = new HashSet<>();
                    dataSourceMap.values().forEach(dataSource -> {
                        Set<String> tabs = DatabaseUtils.queryTables(dataSource, option.getTable());
                        logicTables.addAll(tabs);
                    });
                    if (logicTables.isEmpty()) {
                        logicTables.add(ShardingUtils.getTableTime(option.getTable(), System.currentTimeMillis(), properties.getFormat()));
                        logicTables.add(option.getTable());
                    }
                    String tableStr = logicTables.stream().map(table -> String.format("%s.%s", option.getDatabase(), table)).collect(Collectors.joining(","));
                    ShardingTableRuleConfiguration rule = new ShardingTableRuleConfiguration(option.getTable(), tableStr);
                    final String algorithmName = option.getTable() + "_time";
                    rule.setTableShardingStrategy(new StandardShardingStrategyConfiguration(option.getShardingColumn(), algorithmName));
                    if (StringUtils.isNotEmpty(option.getPrimaryKey()))
                        rule.setKeyGenerateStrategy(new KeyGenerateStrategyConfiguration(option.getPrimaryKey(), "snowflake"));
                    configuration.getTables().add(rule);
                    configuration.getShardingAlgorithms().put(algorithmName, new AlgorithmConfiguration("CLASS_BASED", new Properties() {{
                        put("format", properties.getFormat());
                        put("strategy", "STANDARD");
                        put("algorithmClassName", "cn.hiapi.core.basic.sharding.algorithm.DataMonthShardingAlgorithm");
                    }}));
                } else throw new RuntimeException("配置错误");
            } catch (Exception ex) {
                throw new RuntimeException(ex);
            }
        });
    }


    @Bean
    public DataSource dataSource() throws Exception {
        Map<String, DataSource> dataSourceMap = initDataSource();
        ModeConfiguration modeConfiguration = new ModeConfiguration("Standalone", new StandalonePersistRepositoryConfiguration("JDBC", new Properties()));
        return ShardingSphereDataSourceFactory.createDataSource(modeConfiguration, dataSourceMap,
                Lists.newArrayList(shardingRuleConfiguration(dataSourceMap)),
                new Properties() {{
                    put("sql-show", true);
                }});
    }


}
